"""
第一边和第三边测试
"""
from vision_process.image_deal import *


def find_third_edge_line(totalBoard, totalBoardThres):
    try:
        totalBoardColor = cv.cvtColor(totalBoard, cv.COLOR_GRAY2BGR)
        image_show("TotalBoardThres", totalBoardThres)
        sliceHeight = int(0.8 * totalBoard.shape[0])
        board = totalBoardThres.copy()[sliceHeight:, :]
        scharrX = cv.Scharr(board, cv.CV_64F, 1, 0)
        scharrY = cv.Scharr(board, cv.CV_64F, 0, 1)
        scharrX = cv.convertScaleAbs(scharrX)
        scharrY = cv.convertScaleAbs(scharrY)
        scharrXY = cv.addWeighted(scharrX, 0.1, scharrY, 0.2, 0)
        board = scharrXY.copy()
        image_show("SharrXYThirdEdge", board)
        minLineLengthEdge = int(20 * cp.mm2pix)
        maxLineGap = 5
        yMin = board.shape[0] - int(50 * cp.mm2pix)
        lines = cv.HoughLinesP(board, 1, np.pi / 180, 100, minLineLengthEdge, maxLineGap)
        board1 = totalBoardColor.copy()
        linesList = []
        for line in lines:
            for x1, y1, x2, y2 in line:
                lineLength = distance((x1, y1), (x2, y2))
                if abs(y2 - y1) < 20 * cp.mm2pix and abs(x2 - x1) > 10 * cp.mm2pix and y1 > yMin and \
                        y2 > yMin and lineLength > minLineLengthEdge:
                    y1 += sliceHeight
                    y2 += sliceHeight
                    lineK = round(get_line_k((x1,y1),(x2,y2)),2)
                    linesList.append([x1, y1, (y1 + y2) / 2, lineK,lineLength])
                    linesList.append([x2, y2, (y1 + y2) / 2, lineK,lineLength])
                    cv.line(board1, (x1, y1), (x2, y2), (0, 255, 0), 3)

        # 先去按照lineLength去排序,找到最大的lineLength,如果有超过1/3 物料宽度的,就直接按照这条线去搞定就行了,不用去看其他的线条了.
        linesSortedLength = sorted(linesList,key=lambda x:x[4],reverse=True)
        if len(linesSortedLength) > 0:
            maxLength = int(linesSortedLength[0][4])
            lengthDifToMax = int(maxLength - 0.9 * totalBoard.shape[1])
            print("maxLength = {},lengthDifToMax = {}".format(maxLength,lengthDifToMax))
            if lengthDifToMax >= 0:
                p1 = (linesSortedLength[0][0],linesSortedLength[0][1])
                p2 = (linesSortedLength[1][0],linesSortedLength[1][1])
                bottomCutLineK,bottomCutLineB = get_line_k_and_b(p1,p2)
                return bottomCutLineK,bottomCutLineB
        else:
            return None,None

        # 先按照Y进行排序,从大到小的进行排序
        linesSortedY = sorted(linesList, key=lambda x: x[2], reverse=True)
        print("linesSortedY: {}".format(linesSortedY))
        # 首先是判断方向,看看k为正的个数以及k为负数的个数,k为0的个数,以及k为正数的个数
        # 过滤掉无效的点位,哪些是无效的点位呢,就是那种x坐标范围已经在列表中的点位就是无效点位,怎么样去确定区间的问题,就是要去遍历了.
        if len(linesSortedY) >= 4:
            kList = [x[3] for x in linesSortedY]
            kPosCount = [k for k in kList if k > 0]
            kNegCount = [k for k in kList if k < 0]
            kZeroCount = [k for k in kList if k == 0]
            maxCount = max([kPosCount,kNegCount,kZeroCount])
            if maxCount == kPosCount:
                direction = "pos"
            elif maxCount == kNegCount:
                direction = "neg"
            else:
                direction = "zero"
            if direction != "zero":
                removeSomeWrongPointFlag = True
            else:
                removeSomeWrongPointFlag = False
            print("Direction: {}".format(direction))
            while removeSomeWrongPointFlag:
                for i in range(len(linesSortedY) - 1):
                    xDif = linesSortedY[i + 1][0] - linesSortedY[i][0]
                    if direction == "pos" and xDif > 0:
                        linesSortedY.remove(linesSortedY[i + 1])
                        break
                    if direction == "neg" and xDif < 0:
                        linesSortedY.remove(linesSortedY[i + 1])
                        break
                else:  # 退出循环
                    removeSomeWrongPointFlag = False
        linesList = linesSortedY.copy()
        image_show("thirdEdgeFinded", board1)
        print("linesFinalFinded:{}个,{}".format(len(linesList), linesList))
        if len(linesList) >= 10:
            lineUseStart = 2
            linesRate = 0.65
            linePointsUsedCount = int(len(linesList) * linesRate)  # 选取两个点进行拟合
            linesListFinal = linesList[lineUseStart:lineUseStart + linePointsUsedCount]
            print("linesFinalFinded     :{}".format(linesListFinal))
            linesListFinal = [line[:2] for line in linesListFinal]
            bottomCutLineK, bottomCutLineB = get_linear_fit_k_and_b_horizontal(np.array(linesListFinal))
            bottomCutAngle = np.arctan(bottomCutLineK) * 180 / np.pi
            print("bottomCutAngle: {}".format(bottomCutAngle))
            return bottomCutLineK, bottomCutLineB
        elif len(linesList) >= 4:
            lineUseStart = 1
            linesRate = 0.8
            linePointsUsedCount = int(len(linesList) * linesRate)  # 选取两个点进行拟合
            linePointsUsedCount = max([linePointsUsedCount, 4])  # 就是要保证有2个数以上
            linesListFinal = sorted(linesList, key=lambda x: x[2])[lineUseStart:lineUseStart + linePointsUsedCount]
            print("linesFinalFinded     :{}".format(linesListFinal))
            linesListFinal = [line[:2] for line in linesListFinal]
            bottomCutLineK, bottomCutLineB = get_linear_fit_k_and_b_horizontal(np.array(linesListFinal))
            bottomCutAngle = np.arctan(bottomCutLineK) * 180 / np.pi
            print("bottomCutAngle: {}".format(bottomCutAngle))
            return bottomCutLineK, bottomCutLineB
        else:
            return None, None
    except Exception as e:
        errorInfo = "in find_third_edge_line(): {}".format(str(e))
        print(errorInfo)
        return None, None


def find_first_edge_line(totalBoard, totalBoardThres):
    try:
        totalBoardColor = cv.cvtColor(totalBoard, cv.COLOR_GRAY2BGR)
        image_show("TotalBoardThres", totalBoardThres)
        board = totalBoardThres.copy()[:int(totalBoardThres.shape[0] / 4), :]
        scharrX = cv.Scharr(board, cv.CV_64F, 1, 0)
        scharrY = cv.Scharr(board, cv.CV_64F, 0, 1)
        scharrX = cv.convertScaleAbs(scharrX)
        scharrY = cv.convertScaleAbs(scharrY)
        scharrXY = cv.addWeighted(scharrX, 0.1, scharrY, 0.2, 0)
        board = scharrXY.copy()
        image_show("ScharrXYEdge ", board)
        minLineLengthEdge = int(20 * cp.mm2pix)
        maxLineGap = 5
        yMax = 50 * cp.mm2pix
        lines = cv.HoughLinesP(board, 1, np.pi / 180, 100, minLineLengthEdge, maxLineGap)
        board1 = totalBoardColor.copy()
        linesList = []
        for line in lines:
            for x1, y1, x2, y2 in line:
                lineLength = distance((x1, y1), (x2, y2))
                if abs(y2 - y1) < 20 * cp.mm2pix and abs(
                        x2 - x1) > 10 * cp.mm2pix and y1 < yMax and y2 < yMax and lineLength > minLineLengthEdge:
                    lineK = round(get_line_k((x1,y1),(x2,y2)),2)
                    linesList.append([x1, y1, (y1 + y2) / 2,lineK,lineLength])
                    linesList.append([x2, y2, (y1 + y2) / 2,lineK,lineLength])
                    cv.line(board1, (x1, y1), (x2, y2), (0, 255, 0), 3)
        image_show("EdgeFinded", board1)

        # 先去按照lineLength去排序,找到最大的lineLength,如果有超过1/3 物料宽度的,就直接按照这条线去搞定就行了,不用去看其他的线条了.
        linesSortedLength = sorted(linesList, key=lambda x: x[4], reverse=True)
        if len(linesSortedLength) > 0:
            maxLength = int(linesSortedLength[0][4])
            lengthDifToMax = int(maxLength - 0.25 * totalBoard.shape[1])
            print("maxLength = {},lengthDifToMax = {}".format(maxLength, lengthDifToMax))
            if lengthDifToMax >= 0:
                p1 = (linesSortedLength[0][0], linesSortedLength[0][1])
                p2 = (linesSortedLength[1][0], linesSortedLength[1][1])
                bottomCutLineK, bottomCutLineB = get_line_k_and_b(p1, p2)
                return bottomCutLineK, bottomCutLineB
        else:
            return None, None

        # 先按照Y进行排序
        linesSortedY = sorted(linesList, key=lambda x: x[2])
        if len(linesSortedY) >= 4:
            kList = [x[3] for x in linesSortedY]
            kPosCount = [k for k in kList if k > 0]
            kNegCount = [k for k in kList if k < 0]

            if kPosCount >= kNegCount:
                direction = "pos"
            else:
                direction = "neg"
            # 获取x的方向,正负问题
            removeSomeWrongPointFlag = True
            while removeSomeWrongPointFlag:
                for i in range(len(linesSortedY) - 1):
                    xDif = linesSortedY[i + 1][0] - linesSortedY[i][0]
                    if direction == "pos" and xDif <= 0:
                        linesSortedY.remove(linesSortedY[i + 1])
                        break

                    if direction == "neg" and xDif >= 0:
                        linesSortedY.remove(linesSortedY[i + 1])
                        break
                else:  # 退出循环
                    removeSomeWrongPointFlag = False
        if len(linesList) >= 10:
            lineUseStart = 2
            linesRate = 0.65
            linePointsUsedCount = int(len(linesList) * linesRate)  # 选取最小的六条线进行拟合
            linesListFinal = sorted(linesList, key=lambda x: x[2])[lineUseStart:lineUseStart + linePointsUsedCount]
            print("linesFinalFinded 个数: {} :{}".format(len(linesListFinal), linesListFinal))
            linesListFinal = [line[:2] for line in linesListFinal]
            topCutLineK, topCutLineB = get_linear_fit_k_and_b_horizontal(np.array(linesListFinal))
            topCutAngle = np.arctan(topCutLineK) * 180 / np.pi
            print(topCutAngle)
            return topCutLineK, topCutLineB
        # 取中间的三个点进行拟合
        elif len(linesList) >= 4:
            lineUseStart = 1
            linesRate = 0.8
            linePointsUsedCount = int(len(linesList) * linesRate)  # 选取两个点进行拟合
            linePointsUsedCount = max([linePointsUsedCount, 4])  # 就是要保证有2个数以上
            linesListFinal = sorted(linesList, key=lambda x: x[2])[lineUseStart:lineUseStart + linePointsUsedCount]
            print("linesFinalFinded 个数: {} :{}".format(len(linesListFinal), linesListFinal))
            linesListFinal = [line[:2] for line in linesListFinal]
            topCutLineK, topCutLineB = get_linear_fit_k_and_b_horizontal(np.array(linesListFinal))
            topCutAngle = np.arctan(topCutLineK) * 180 / np.pi
            print(topCutAngle)
            return topCutLineK, topCutLineB
        else:
            return None, None
    except Exception as e:
        print(str(e))
        return None, None


def get_board_edge(imageTest, totalBoardBox, boardThres):
    totalBoard = get_board_widthout_angle(imageTest, totalBoardBox)
    totalBoardThresWithoutAngle = get_board_widthout_angle(boardThres, totalBoardBox)

    leftTop = (0, 0)
    rightTop = (totalBoard.shape[1], 0)
    rightBottom = (totalBoard.shape[1], totalBoard.shape[0])
    leftBottom = (0, totalBoard.shape[0])

    topCutLineK, topCutLineB = find_first_edge_line(totalBoard, totalBoardThresWithoutAngle)
    bottomLineK, bottomLineB = find_third_edge_line(totalBoard, totalBoardThresWithoutAngle)
    leftCutLineK, leftCutLineB = get_line_k_and_b(leftTop, leftBottom)
    rightCutLineK, rightCutLineB = get_line_k_and_b(rightTop, rightBottom)

    if topCutLineK is None:
        topCutLineK, topCutLineB = get_line_k_and_b(leftTop, rightTop)

    if bottomLineK is None:
        bottomLineK, bottomLineB = get_line_k_and_b(leftBottom, rightBottom)

    newLeftTop = get_two_lines_intersection(topCutLineK, topCutLineB, leftCutLineK, leftCutLineB)
    newRightTop = get_two_lines_intersection(topCutLineK, topCutLineB, rightCutLineK, rightCutLineB)
    newLeftBottom = get_two_lines_intersection(leftCutLineK, leftCutLineB, bottomLineK, bottomLineB)
    newRightBottom = get_two_lines_intersection(rightCutLineK, rightCutLineB, bottomLineK, bottomLineB)
    cv.destroyAllWindows()
    totalBoardExtend = get_board_widthout_angle(imageTest, totalBoardBox, extend=True)
    image_show("TotalBoardExtend", totalBoardExtend)
    totalBoardColor = cv.cvtColor(totalBoardExtend, cv.COLOR_GRAY2BGR)
    pList = [(int(p[0]) + 50, int(p[1]) + 50) for p in [newLeftTop, newRightTop, newRightBottom, newLeftBottom]]
    for index, p in enumerate(pList):
        cv.line(totalBoardColor, p, pList[(index + 1) % 4], (0, 255, 0), 3)
    image_show("TotalBoardRes", totalBoardColor)
    return totalBoardColor


if __name__ == '__main__':
    # 先加载参数
    filePath = r"D:\xizhen_01\first_edge_find"
    isFirstBoard = False  # 是否是首片,首片有首片的算法逻辑
    res, data = db.get_line_cam_check_thres()
    totalHeightMax = data.get("totalHeightMax")
    totalHeightMin = data.get("totalHeightMin")
    boardBreak = False
    calcThresBig = 20
    calcThresSmall = 9
    innerToEdgeThres = 5
    mtx = [
        [0, 1, 0],
        [1, 0, 0]
    ]
    remainEdge = [0, 21, 0, 21]
    topRemain, leftRemain, bottomRemain, rightRemain = [int(x * cp.mm2pix) for x in remainEdge]
    keepDir = filePath
    if not os.path.isdir(keepDir):
        keepDir, _ = os.path.split(filePath)

    keepOkPath = os.path.join(keepDir, "ok")
    keepOkResPath = os.path.join(keepDir, "total_board_res")
    if not os.path.exists(keepOkPath):
        os.mkdir(keepOkPath)
    if not os.path.exists(keepOkResPath):
        os.mkdir(keepOkResPath)

    fileNames = []  # 保存要去识别的所有的图片
    if os.path.isfile(filePath):
        filePath, fileName = os.path.split(filePath)
        fileNames.append(fileName)
    else:
        for root, dirs, files in os.walk(filePath):
            if root != filePath:
                continue
            for file in files:
                if file[-4:] != ".bmp":
                    continue
                fileNames.append(file)

    for index, fileName in enumerate(fileNames):
        filePathReal = os.path.join(filePath, fileName)
        srcW, srcH, upLineOffset, downLineOffset = get_width_height_by_image_name(fileName)
        if srcW is None or srcH is None:
            print("*" * 10 + "第 {} 张图 ({}),图像名称错误,已经跳过".format(index + 1, fileName) + "*" * 10)
            continue
        if 0 in upLineOffset or 0 in downLineOffset:
            isFirstBoard = True
        else:
            isFirstBoard = False
        print("*" * 10 + " 第 {} 张图 ({})".format(index + 1, fileName) + "*" * 10)
        imageTest = cv.imread(filePathReal, cv.IMREAD_GRAYSCALE)
        calcStart = time.time()
        checkRes, checkInfo, totalBoardBox, boardThres, totalBoardRes, mm = check_total_board_is_ok(imageTest, srcW, srcH,
                                                                                                    boardBreak,
                                                                                                    totalHeightMax,
                                                                                                    totalHeightMin)

        totalBoardEdgeRes = get_board_edge(imageTest, totalBoardBox, boardThres)
        print("算法耗时: {}".format(time.time() - calcStart))
        os.remove(filePathReal)
        imageRes = get_size_scale_image(totalBoardEdgeRes, sizeScale=4)
        cv.imwrite(os.path.join(keepOkResPath, fileName), imageRes)
        cv.imwrite(os.path.join(keepOkPath, fileName), imageTest)
        print("识别成功")
        continue
